import numpy as np
import pandas as pd
from environment import Plane


Timetable = pd.read_csv("F:/plane_information/timetable.csv" , nrows = 10 , usecols = [0 , 1 , 2 , 3 , 4])   # 改nrows的个数为需要的飞机个数

class RL(object):

    #初始化
    def __init__(self , actions_space , learning_rate = 0.9, reward_decay = 0.9 , e_greedy = 0.9 , step = 2):
        self.actions = actions_space
        self.lr = learning_rate
        self.gamma = reward_decay
        self.epsilon = e_greedy     #选择性修改：ε-greedy的选择
        self.sarsa_table = pd.DataFrame(columns = self.actions , dtype = np.float64)
        self.step = step

    '''
    # 选择动作
    def choose_action(self , observation , max_delay_plane):

        # 检查是否在Sarsa_table中
        self.check_and_update(observation)

        #找状态对应的动作
        actions = self.sarsa_table[0].tolist()
        observation_actions = self.sarsa_table.loc[actions.index(observation) , 1 : ]

        # 90%的可能性找里面动作最大的，10%随便找一个
        if np.random.rand() < self.epsilon:
            action = np.random.choice(observation_actions[observation_actions == np.max(observation_actions)].index)
        else:
            action = np.random.choice(observation_actions.index)

        #拿到奖励值
        reward = observation_actions[action]

        #拿到对应的下一个状态
        return action , reward
        '''

    # 2.2 找到奖励值最大的动作,跳转到下一个状态,并记录该状态下的延误值
    def choose_action(self , observation):
        pd.set_option('display.width', None)
        # 2.2.1 找到该状态下的延误值最小的动作，并拿到这个数据
        actions = self.sarsa_table[0].tolist()
        observation_actions = self.sarsa_table.loc[actions.index(observation), 1:]
        best_action = np.random.choice(observation_actions[observation_actions == np.min(observation_actions)].index)
        dtime = np.min(observation_actions)

        # 2.2.2 如果说找了一圈动作没有你的小，则选择"原地不动：动作"，并拿到下一个状态值

        if best_action == 21:
            observation = observation.copy()
        else:
            t_1 = observation.index(best_action) - 1
            t = observation.index(best_action)
            observation[t_1] , observation[t] = observation[t] , observation[t_1]

        return dtime , observation

    #检查并更新 Sarsa表
    def check_and_update(self , state):
        if state not in self.sarsa_table[0].tolist(): #改：新建sarsa_table的新建，对应飞机的动作，有几个写到几个。
            self.sarsa_table = self.sarsa_table.append({0 : state , 1 : 0.0 , 2 : 0.0 , 3 : 0.0 , 4 : 0.0 , 5 : 0.0 , 6:0.0 , 7:0.0 , 8:0.0 , 9:0.0 , 10:0.0 , 11:0.0 } , ignore_index = True)

class SarsaTable(RL):

    #继承 RL模块初始化
    def __init__(self , actions , learning_rate = 0.9 , reward_decay = 0.9 , e_greedy = 0.9 , step = 2):
        super(SarsaTable , self).__init__(actions , learning_rate , reward_decay , e_greedy , step)

    # 2.1.2 拿到延误数据，写入Q表
    def writing_in_sarsa_table(self , action, dtime, ob):
        self.check_and_update(ob)
        all_observation = self.sarsa_table[0].tolist()
        index = all_observation.index(ob)
        self.sarsa_table.loc[index, action] = dtime

    # 2.3 根据动作重新更新一次Q表的策略
    def learn(self , observation, action, dtime, next_observation):
        self.check_and_update(next_observation)
        all_observation = self.sarsa_table[0].tolist()
        now_index = all_observation.index(observation)
        index = all_observation.index(next_observation)
        sarsa_predict = self.sarsa_table.loc[now_index , action]
        sarsa_target = dtime + self.gamma * self.sarsa_table.loc[index , 1:].min()
        self.sarsa_table.loc[now_index , action] += self.lr * (sarsa_target - sarsa_predict)

'''#学习
    def learn(self, s, a, r, next_s):
        # 检查是否在Q表中，不在就加进去,并且得到最终标记 terminal个数
        self.check_and_update(next_s)
        self.check_and_update(s)

        terminal = len(s) - 1
        a += 1
        # 找到对应的下标
        plane_row = self.sarsa_table[0].tolist()
        s_index = plane_row.index(s)
        next_s_index = plane_row.index(next_s)
        sarsa_predict = self.sarsa_table.loc[s_index , a]

        sarsa_target = r

        #更新数据
        self.sarsa_table.loc[s_index , a] += self.lr * (sarsa_target - sarsa_predict)
        pd.set_option('display.width', 1000)
        '''
